#!/home/robot/anaconda3/envs/yolograsp/bin/python3
import rospy
import copy
from vision_messages.srv import GraspPoseEst, GraspPoseEstResponse
import time
import ros_np_multiarray as ros_nm
import numpy as np
import sys, os, yaml
import argparse
############## Import Local Module ############
sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), 'lib'))
from lib.shape_encoder.demo import load_grasp_source, load_model_encoder, load_model_dif
from lib.shape_encoder.demo import transfer_grasp, generate_mug_body_grasp
from lib.grasp_refine.demo import refine_grasp
# sys.path.pop(0)
###############################################

def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--grasp_choose_cfg', type=str, default='')
    parser.add_argument("--icp", action='store_true', default=False)
    opt = parser.parse_args()
    return opt
    
def correct_grasp_pose(RT, category, cfg):
    """
    Make sure that 
    1. one-finger side is closer to the obj center for mug and bowl.
    2. camera is up when graspping bottle.  
    """
    if (category == 'mug' and cfg['type'] == 'rim') or category == 'bowl':
        point1 = [-7.27595772e-12, -0.04, 0.065962+0.053717]
        point2 = [-7.27595772e-12, 0.04, 0.065962+0.053717]
        R = RT[:3, :3]
        t = RT[:3, 3]
        # transpose these points to the obj center
        point1 = R @ point1 + t
        point2 = R @ point2 + t
        # get the distance of them to the obj center
        dist1 = np.linalg.norm(point1)
        dist2 = np.linalg.norm(point2)
        if dist1 > dist2:
            # rotate RT along z axis by 180 degree
            delta_RT = np.eye(4)
            delta_RT[:3, :3] = np.array([[-1, 0, 0], [0, -1, 0], [0, 0, 1]])
            RT = RT @ delta_RT
        panda_2_jaco = np.array([[0,1,0],[-1,0,0],[0,0,1]], dtype=np.float32)
        RT[:3,:3] = np.matmul(RT[:3,:3], panda_2_jaco)
    elif (category == 'mug' and cfg['type'] == 'handle') or category == 'bottle':
        panda_2_jaco = np.array([[0,1,0],[-1,0,0],[0,0,1]], dtype=np.float32)
        RT[:3,:3] = np.matmul(RT[:3,:3], panda_2_jaco)
        grasp_pose_y = RT[:3, 1]
        flag = np.sum(grasp_pose_y * np.array([0,1,0]))
        if flag < 0:
            # rotate RT along z axis by 180 degree
            delta_RT = np.eye(4)
            delta_RT[:3, :3] = np.array([[-1, 0, 0], [0, -1, 0], [0, 0, 1]])
            RT = RT @ delta_RT
    elif category == 'mug' and cfg['type'] == 'body':
        panda_2_jaco = np.array([[0,1,0],[-1,0,0],[0,0,1]], dtype=np.float32)
        RT[:3,:3] = np.matmul(RT[:3,:3], panda_2_jaco)
        grasp_pose_y = RT[:3, 1]
        flag = np.sum(grasp_pose_y * np.array([0,1,0]))
        if flag < 0:
            # rotate RT along z axis by 180 degree
            delta_RT = np.eye(4)
            delta_RT[:3, :3] = np.array([[-1, 0, 0], [0, -1, 0], [0, 0, 1]])
            RT = RT @ delta_RT
    return RT

def get_RT_from_grasp_params(panda_grasp_info, obj_scale, category, cfg):
    left_points = panda_grasp_info['left_points']
    right_points = panda_grasp_info['right_points']
    appro = panda_grasp_info['approach_vector']
    depth = panda_grasp_info['depth']

    hori = right_points - left_points
    hori = hori / np.linalg.norm(hori)
    normal = np.cross(hori, appro)
    normal = normal / np.linalg.norm(normal)

    appro_correct = np.cross(normal, hori)

    RT = np.eye(4, dtype=np.float32)
    RT[:3,:3] = np.array([normal, hori, appro_correct], dtype=np.float32).T
    RT[:3,3] = (left_points + right_points) *  (obj_scale / 2) / 2 

    RT = correct_grasp_pose(RT, category, cfg)

    return RT

class GraspEstServer():
    def __init__(self, icp, grasp_choose_cfg):
        self.icp=icp
        if self.icp:
            rospy.loginfo('  Using ICP ')
        # model and grasp source
        grasp_source_root = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'panda_grasp_data')
        self.refine_pth_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'refine_grasp')
        lib_root = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'lib')
        cate_list = ['mug', 'bowl', 'bottle']
        self.model_encoder = dict()
        self.model_dif = dict()
        self.grasp_source = dict()
        self.grasp_work = grasp_choose_cfg.replace('_grasp.yaml', '')
        with open(os.path.join(grasp_source_root, grasp_choose_cfg), 'r') as stream:
            self.grasp_src_cfg = yaml.safe_load(stream)
        for cate in cate_list:
            self.model_encoder[cate] = load_model_encoder(lib_root, cate)
            dif_config = '{}/dif/configs/generate/{}.yml'.format(lib_root, cate)
            with open(os.path.join(dif_config),'r') as stream: 
                meta_params = yaml.safe_load(stream)
            self.model_dif[cate] = load_model_dif(lib_root, meta_params)
            if cate not in self.grasp_src_cfg.keys():
                raise ValueError('There is no \'{}\' config in file {}'.format(cate, grasp_choose_cfg))
            self.grasp_source[cate] = load_grasp_source(grasp_source_root, cate, self.grasp_src_cfg[cate])
        # ros node and service
        rospy.init_node('grasp_est_server')
        rospy.Service('/perception/yolograsp/grasp_est', GraspPoseEst, self.grasp_est_callback)
        rospy.loginfo('grasp_est_server init success!')

    def grasp_est_callback(self, req):
        scale = req.single_obj_pose.scale
        cate = req.single_obj_pose.category_name
        obj_pcs = req.single_obj_pose.obj_pcs
        obj_pcs = ros_nm.to_numpy_f32(obj_pcs)

        pose_results = {'pred_s': scale, 'obj_pcs': obj_pcs, 'category': cate}

        # for icp
        cam_pcs = req.single_obj_pose.cam_pcs
        cam_pcs = ros_nm.to_numpy_f32(cam_pcs)
        pose_results['cam_pcs'] = cam_pcs        
        RT = req.single_obj_pose.RT
        RT = ros_nm.to_numpy_f32(RT)
        sRT = RT
        sRT[:3,:3] *= scale
        pose_results['sRT'] = sRT

        start_time = time.time()
        
        if self.grasp_work == 'pour_water' and self.grasp_src_cfg['mug']['type'] == 'body':
            panda_grasp, pose_icp = generate_mug_body_grasp(pose_results, self.model_encoder['mug'], self.model_dif['mug'], self.grasp_source['mug'],
                                                icp=self.icp)
            trasfer_time = time.time()
            rospy.loginfo('trasfer grasp use time: {}'.format(trasfer_time - start_time))
            pose_icp_response = req.single_obj_pose                                                
            if self.icp and cate != 'bowl':
                RT_icp = ros_nm.to_multiarray_f32(pose_icp['RT_icp'])
                pose_icp_response.RT = RT_icp
                pose_icp_response.scale = pose_icp['pred_s']
                
        else:
            panda_grasp, pose_icp = transfer_grasp(pose_results, self.model_encoder[cate], self.model_dif[cate], self.grasp_source[cate],
                                                icp=self.icp)

            trasfer_time = time.time()
            rospy.loginfo('trasfer grasp use time: {}'.format(trasfer_time - start_time))
            pose_icp_response = req.single_obj_pose

            if self.icp and cate != 'bowl':
                RT_icp = ros_nm.to_multiarray_f32(pose_icp['RT_icp'])
                pose_icp_response.RT = RT_icp
                pose_icp_response.scale = pose_icp['pred_s']

            panda_grasp = refine_grasp(panda_grasp, self.model_dif[cate], refine_pth_dir=self.refine_pth_dir)
            refine_time = time.time()
            rospy.loginfo('refine grasp use time: {}'.format(refine_time - trasfer_time))

        assert isinstance(panda_grasp['grasp_params'], dict)

        grasp_pose = get_RT_from_grasp_params(panda_grasp['grasp_params'], scale, cate, self.grasp_src_cfg[cate])
        grasp_pose = ros_nm.to_multiarray_f32(grasp_pose)

        rospy.loginfo('Process Done! Return Grasp and Object Pose Results!')

        return GraspPoseEstResponse(grasp_pose, pose_icp_response)


if __name__ == '__main__':
    args = get_args()
    grasp_est_server = GraspEstServer(icp=args.icp, grasp_choose_cfg=args.grasp_choose_cfg)
    rospy.spin()
